import os
import psycopg2
from psycopg2 import sql

# 配置数据库连接
DB_CONFIG = {
    'dbname': 'postgres',
    'user': 'wx',
    'password': '123456',
    'host': 'localhost',  # 或你的数据库地址
}

# 存放 SQL 文件路径的记录文件
LOG_FILE = "executed.log"


def get_executed_files(log_file):
    """读取已执行的 SQL 文件列表"""
    if os.path.exists(log_file):
        with open(log_file, "r") as f:
            return set(f.read().splitlines())
    return set()


def save_executed_file(log_file, file_name):
    """记录已执行的 SQL 文件"""
    with open(log_file, "a") as f:
        f.write(file_name + "\n")


def execute_sql_file(cursor, file_path):
    """执行单个 SQL 文件"""
    with open(file_path, "r") as f:
        sql_content = f.read()
    cursor.execute(sql.SQL(sql_content))


def run_sql_files(directory, log_file):
    """运行目录下所有未执行过的 SQL 文件"""
    executed_files = get_executed_files(log_file)
    all_files = [f for f in os.listdir(directory) if f.endswith(".sql")]

    # 过滤出未执行的 SQL 文件
    new_files = [f for f in all_files if f not in executed_files]

    if not new_files:
        print("没有新的 SQL 文件需要执行。")
        return

    try:
        # 连接到 PostgreSQL 数据库
        with psycopg2.connect(**DB_CONFIG) as conn:
            with conn.cursor() as cursor:
                for file_name in sorted(new_files):  # 按文件名排序执行
                    file_path = os.path.join(directory, file_name)
                    print(f"正在执行文件: {file_name}")
                    execute_sql_file(cursor, file_path)
                    save_executed_file(log_file, file_name)
                    conn.commit()  # 提交事务
                    print(f"文件 {file_name} 执行完成！")
    except Exception as e:
        print(f"执行 SQL 文件时发生错误: {e}")
    else:
        print("所有新 SQL 文件已成功执行！")


if __name__ == "__main__":
    # 指定存放 SQL 文件的目录
    SQL_DIRECTORY = "./wx"  # 修改为你的实际目录

    # 确保目录存在
    if not os.path.exists(SQL_DIRECTORY):
        print(f"指定的目录不存在: {SQL_DIRECTORY}")
    else:
        run_sql_files(SQL_DIRECTORY, LOG_FILE)

